import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense,Dropout
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.preprocessing import LabelEncoder
from tensorflow.keras.utils import to_categorical
import numpy as np
import random
# 设置随机数种子
seed = 42
np.random.seed(seed)
tf.random.set_seed(seed)
random.seed(seed)

def LSTMForSmooth(dataX,dataY):
    # 构建 LSTM 模型
    model = Sequential()
    model.add(LSTM(units=32, input_shape=(dataX.shape[-2], dataX.shape[-1]), return_sequences=True, activation='relu'))
    model.add(LSTM(units=16, return_sequences=False, activation='relu'))
    model.add(Dense(dataY.shape[-1]))
    # 编译模型
    model.compile(optimizer='adam', loss='mse')
    # 训练模型
    model.fit(dataX, dataY, epochs=100, batch_size=20)
    y_pred = model.predict(dataX)
    return y_pred


def LSTMForClassify(dataX,labels):
    # 将标签进行one-hot编码
    encoder = LabelEncoder()
    labels = encoder.fit_transform(labels)
    labels = to_categorical(labels, num_classes=np.max(labels)+1)

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(dataX, labels, test_size=0.3, random_state=seed)

    # 创建LSTM模型
    model = Sequential()
    model.add(LSTM(32, input_shape=(X_train.shape[-2], X_train.shape[-1]),return_sequences=True, activation='relu'))
    model.add(LSTM(units=16, return_sequences=False, activation='relu'))
    model.add(Dropout(0.2))
    model.add(Dense(labels.shape[-1], activation='softmax'))

    # 编译模型
    model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

    # 训练模型
    model.fit(X_train, y_train, epochs=100, batch_size=5)

    # 评估模型
    loss, accuracy = model.evaluate(X_test, y_test)
    print(f'Test Accuracy: {accuracy}')



